import torch
class BatchIter:
    def __init__(self, loader, cuda):
        self.batch_iterator = iter(loader)
        self.batch_size = loader.batch_size
        self.loader = loader
        self.cuda = cuda

    def next(self):
        try:
            data = next(self.batch_iterator)
        except StopIteration:
            self.batch_iterator = iter(self.loader)
            data = next(self.batch_iterator)

        if len(data) == 2:
            images, targets = data
        elif len(data) == 3:
            images, targets, sla = data

        if self.cuda:
            images = images.cuda()
            with torch.no_grad():
                targets = targets.cuda()

            if len(data) == 3:
                sla = sla.cuda()

        if len(data) == 3:
            return images, targets, sla
        else:
            return images, targets